#ifndef HEADER_fd_src_discof_repair_fd_repair_h
#define HEADER_fd_src_discof_repair_fd_repair_h

/* fd_repair implements the Solana Repair protocol.  In a nutshell,
   Repair is a protocol for recovering shreds that the validator is
   expecting but has not received from Turbine.  Neither the logic for
   how the validator determines it should be expecting a shred nor the
   logic for determining which peer validator to request the shred from
   is implemented in this file (see fd_policy.h instead); rather, this
   is an implementation of the protocol itself for requesting shreds
   from other validators.

   The repair protocol supports four different message types on the
   client side:

   - Pong( ping_token )

     This is a response message to address validation via a Ping-Pong
     protocol.  When a validator receives a repair request from another
     validator it does not recognize, it will ignore the request and
     instead respond with its own Ping message to the requesting
     validator.  The Ping contains a 32-byte token that the requesting
     validator needs to hash and sign as part of a Pong response.  The
     scheme is as follows: the Ping token is concatenated with the
     prefix "SOLANA_PING_PONG", which is then piped in as the preimage
     into SHA-256.  The resulting hash is then signed and the signature
     along with the actual hash is packed into a Pong.

   - Shred( slot, shred_idx )

     This is a request for a specific shred in the provided slot.  The
     (slot, shred index) generally uniquely identifies a shred except in
     certain exceptional conditions (equivocation).  The responding
     validator will return the shred if it has it.

   - HighestShred( slot, shred_idx )

     This is a request for the highest shred in the provided slot that
     is greater than or equal to shred_idx.  Note this is not
     necessarily the last shred in the slot, as it depends on what the
     responding validator has available.  The responding validator will
     return the highest shred it has that meets the condition.

   - Orphan( slot )

     This is a request for up to 10 shreds, where each shred is the
     prior one's ancestor, beginning from but excluding slot.  For
     example, an orphan request for slot 10 will return a single shred
     for slots 9, 8, 7, 6, 5, 4, 3, 2 and 1 (assuming no skips).  Also,
     the responding validator will return the highest shred index it has
     for every ancestor it knows about.

    All 3 repair request types are prefixed with a common header of from
    pubkey, to pubkey, ulong timestamp and uint nonce.  The timestamp is
    a standard UNIX epoch (milliseconds since 1970-01-01T00:00:00Z) and
    the nonce is echoed back by the responding validator in the repair
    response.  Unlike a typical cryptographic nonce that prevents replay
    attacks, the repair server implementation ignores the nonce and this
    implementation leaves it up to the calling application to manage the
    nonce.  Note the 4 nonce bytes are appended after the end of a shred
    in a repair response.

    All communication across the wire is done with bincode serialization
    encoding.  Serialization and deserialization is code-generated by
    fd_types. */

#include "../../ballet/ed25519/fd_ed25519.h"
#include "../../flamenco/types/fd_types_custom.h"

/* FD_REPAIR_USE_HANDHOLDING:  Define this to non-zero at compile time
   to turn on additional runtime checks and logging. */

#ifndef FD_REPAIR_USE_HANDHOLDING
#define FD_REPAIR_USE_HANDHOLDING 1
#endif

/* FD_REPAIR_KIND_{PONG,SHRED,HIGHEST_SHRED,ORPHAN} specify discriminant
   values the protocol uses to distinguish message types. */

#define FD_REPAIR_KIND_PING          (0U)
#define FD_REPAIR_KIND_PONG          (7U)
#define FD_REPAIR_KIND_SHRED         (8U)
#define FD_REPAIR_KIND_HIGHEST_SHRED (9U)
#define FD_REPAIR_KIND_ORPHAN        (10U)

/* fd_repair_pong describes the schema of a Pong. */

struct __attribute__((packed)) fd_repair_pong {
  fd_pubkey_t      from; /* pubkey of the validator responding with the pong */
  fd_hash_t        hash; /* sha-256 hash generated from a ping hash */
  fd_ed25519_sig_t sig;  /* from's signature over the preceding hash field */
};
typedef struct fd_repair_pong fd_repair_pong_t;

/* REQ_HDR defines the common header of Repair request types. */

#define REQ_HDR                                                                     \
  fd_ed25519_sig_t sig;   /* ed25519 signature over all the subsequent fields */    \
  fd_pubkey_t      from;  /* pubkey of the validator that sent the request */       \
  fd_pubkey_t      to;    /* pubkey of the validator that is being requested */     \
  ulong            ts;    /* timestamp in milliseconds since unix epoch */          \
  uint             nonce; /* nonce to be echoed back by the responding validator */ \

/* fd_repair_shred requests the specific shred at slot and shred_idx
   from a peer validator. */

/* TODO: remove _req suffix from below after we remove all fd_types inclusions in repair_tile */
struct __attribute__((packed)) fd_repair_shred_req {
  REQ_HDR
  ulong slot;
  ulong shred_idx;
};
typedef struct fd_repair_shred_req fd_repair_shred_req_t;

/* fd_repair_highest_shred requests the highest shred in slot greater
   than shred_idx that a peer validator has.  Note this is not
   necessarily the last shred in the slot, as it depends on what the
   peer has available. */

struct __attribute__((packed)) fd_repair_highest_shred_req {
  REQ_HDR
  ulong slot;
  ulong shred_idx;
};
typedef struct fd_repair_highest_shred_req fd_repair_highest_shred_req_t;

/* fd_repair_orphan requests the ancestors of slot (an "orphaned" slot)
   from a peer validator.  The peer can respond with shreds for up to 10
   ancestor slots, where every shred is the last shred for that slot. */

struct __attribute__((packed)) fd_repair_orphan_req {
  REQ_HDR
  ulong slot;
};
typedef struct fd_repair_orphan_req fd_repair_orphan_req_t;

/* fd_repair_msg_t defines the schema of all Repair message types. */

struct __attribute__((packed)) fd_repair_msg {
  uint kind; /* FD_REPAIR_KIND_{PONG,SHRED,HIGHEST_SHRED,ORPHAN} */
  union {
    fd_repair_pong_t              pong;
    fd_repair_shred_req_t         shred;
    fd_repair_highest_shred_req_t highest_shred;
    fd_repair_orphan_req_t        orphan;
  };
};
typedef struct fd_repair_msg fd_repair_msg_t;

struct __attribute__((packed)) fd_repair_ping {
  uint kind;
  fd_repair_pong_t ping;
};
typedef struct fd_repair_ping fd_repair_ping_t;

/* FD_REPAIR_PONG_PREIMAGE_PREFIX is used by Repair's Ping-Pong protocol.
   Both a Ping and Pong contain a hash token, that is generated from a
   preimage prefixed with the below.  */

#define FD_REPAIR_PONG_PREIMAGE_PREFIX "SOLANA_PING_PONG"
#define FD_REPAIR_PONG_PREIMAGE_SZ (48UL)

/* FD_REPAIR_MAX_PREIMAGE_SZ is the maximum size of a preimage for a
   repair request. This is the size of the largest repair request
   (highest_shred or shred) without the signature. */

#define FD_REPAIR_MAX_PREIMAGE_SZ (sizeof(fd_repair_msg_t) - sizeof(fd_ed25519_sig_t))

static const fd_pubkey_t null_pubkey = {{ 0 }};

/* fd_repair_sign_fn defines the function signature for a user-provided
   signing callback. */

typedef void (fd_repair_sign_fn)( void * ctx, fd_repair_msg_t * msg, uchar * sig );

struct fd_repair {
  fd_pubkey_t         identity_key; /* validator identity key */
  fd_repair_sign_fn * sign_fn;      /* user-provided signing callback */
  void *              sign_ctx;     /* user-provided context for signing callback */
  fd_repair_msg_t     msg;          /* buffer for outgoing repair requests */
};
typedef struct fd_repair fd_repair_t;

/* Constructors */

/* fd_repair_{align,footprint} return the required alignment and
   footprint of a memory region suitable for use as repair.  Declaration
   friendly (e.g. a memory region declared as "fd_repair_t repair[1];"
   will automatically have the needed alignment and footprint). */

FD_FN_CONST static inline ulong
fd_repair_align( void ) {
  return alignof(fd_repair_t);
}

FD_FN_CONST static inline ulong
fd_repair_footprint( void ) {
  return sizeof(fd_repair_t);
}

/* fd_repair_new formats an unused memory region for use as a repair.
   mem is a non-NULL pointer to this region in the local address space
   with the required footprint and alignment.
   Initializes repair with the public identity key. */

void *
fd_repair_new( void * shmem, fd_pubkey_t * identity_key );

/* fd_repair_join joins the caller to the repair.  repair points to the
   first byte of the memory region backing the repair in the caller's
   address space.  Returns a pointer in the local address space to
   repair on success. */

fd_repair_t *
fd_repair_join( void * repair );

/* fd_repair_leave leaves a current local join.  Returns a pointer to
   the underlying shared memory region on success and NULL on failure
   (logs details).  Reasons for failure include repair is NULL. */

void *
fd_repair_leave( fd_repair_t const * repair );

/* fd_repair_delete unformats a memory region used as a repair.  Assumes
   only the nobody is joined to the region.  Returns a pointer to the
   underlying shared memory region or NULL if used obviously in error
   (e.g. repair is obviously not a repair ... logs details).  The
   ownership of the memory region is transferred to the caller. */

void *
fd_repair_delete( void * repair );

/* fd_repair_{pong,shred,highest_shred,orphan} creates and returns a
   pointer to a serialized {Pong,Shred,HighestShred,Orphan} message.
   Does not require the caller to provide memory, as Repair itself
   maintains a dedicated memory region (repair->msg) for buffering
   requests.  Assumes repair->msg is not already buffering an existing
   request and can be overwritten.  Returns a pointer to repair->msg on
   success, NULL on failure. */


fd_repair_msg_t * fd_repair_pong         ( fd_repair_t * repair, fd_hash_t * ping_token );
fd_repair_msg_t * fd_repair_shred        ( fd_repair_t * repair, fd_pubkey_t const * to, ulong ts, uint nonce, ulong slot, ulong shred_idx );
fd_repair_msg_t * fd_repair_highest_shred( fd_repair_t * repair, fd_pubkey_t const * to, ulong ts, uint nonce, ulong slot, ulong shred_idx );
fd_repair_msg_t * fd_repair_orphan       ( fd_repair_t * repair, fd_pubkey_t const * to, ulong ts, uint nonce, ulong slot );

/* fd_repair_sz returns the bincode-serialized sz of msg. */

FD_FN_PURE static inline ulong
fd_repair_sz( fd_repair_msg_t const * msg ) {
   switch( msg->kind ) {
     case FD_REPAIR_KIND_PONG:          return sizeof(uint) + sizeof(fd_repair_pong_t);
     case FD_REPAIR_KIND_SHRED:         return sizeof(uint) + sizeof(fd_repair_shred_req_t);
     case FD_REPAIR_KIND_HIGHEST_SHRED: return sizeof(uint) + sizeof(fd_repair_highest_shred_req_t);
     case FD_REPAIR_KIND_ORPHAN:        return sizeof(uint) + sizeof(fd_repair_orphan_req_t);
     default:                           FD_LOG_ERR(( "Unhandled repair kind %u", msg->kind ));
   }
}

/* preimage_pong takes a pong and takes the token in a ping and a
   FD_REPAIR_PONG_PREIMAGE_SZ-byte buffer and returns a pointer to a
   preimage that can be signed. */

static inline uchar *
preimage_pong( fd_hash_t const * ping_token, uchar * preimage_buf, ulong preimage_sz ) {
# if FD_REPAIR_USE_HANDHOLDING
  if( FD_UNLIKELY( preimage_sz != FD_REPAIR_PONG_PREIMAGE_SZ ) ) {
    FD_LOG_ERR(( "preimage_sz %lu must be %lu", preimage_sz, FD_REPAIR_PONG_PREIMAGE_SZ ));
  }
# endif
  ulong prefix_sz = sizeof(FD_REPAIR_PONG_PREIMAGE_PREFIX) - 1 /* subtract NUL */;
  memcpy( preimage_buf,        FD_REPAIR_PONG_PREIMAGE_PREFIX, prefix_sz         );
  memcpy( preimage_buf + prefix_sz, ping_token, sizeof(fd_hash_t) );
  return preimage_buf;
}

/* preimage_req takes a repair request populated with all fields except
   for the signature, and returns a pointer to a preimage that can be
   signed.  Modifies the msg in place.

   At the start of this function, the repair_msg_t should contain

     [ discriminant ] [ empty sig ] [ repair request fields ]
     ^                ^             ^
     0                4             68

     https://github.com/solana-labs/solana/blob/master/core/src/repair/serve_repair.rs#L1258

     We want to sign over
     [ discriminant ] [ payload ]
      ^                ^
      buffer           buffer+4

   We can do this without using any extra memory copying the discriminant
   to the last 4 bytes of the sig field, and returning a pointer to
   that copied location.  The sig field should be overwritten with the
   actual signature value later, else the fd_repair_msg_t will be
   incorrect. */

static inline uchar *
preimage_req( fd_repair_msg_t * msg, ulong * preimage_sz ) {
  uchar * preimage = (uchar *)fd_type_pun(msg);
  preimage += sizeof(fd_ed25519_sig_t);
  FD_STORE( uint, preimage, msg->kind ); /* copy discriminant over */
  *preimage_sz = fd_repair_sz( msg ) - sizeof(fd_ed25519_sig_t);
  return preimage;
}
#endif /* HEADER_fd_src_discof_repair_fd_repair_h */
